// Copyright (c) 2023 by Rockchip Electronics Co., Ltd. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

/*-------------------------------------------
                Includes
-------------------------------------------*/
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#include "cls.h"
#include "cls_process.h"
#include "image_utils.h"
#include "file_utils.h"
#include "image_drawing.h"

#include <iostream> 
#include <fstream>
#include <nlohmann/json.hpp>
using namespace std;


void read_json(string label_path,int &class_num,nlohmann::json &label_json){
    ifstream json_file(label_path);
    json_file >> label_json;
    class_num = static_cast<int>(label_json.size());
    json_file.close();
}


/*-------------------------------------------
                  Main Function
-------------------------------------------*/
int main()
{
    const char *model_path = "../model/cls/head/cls_head.rknn";
    const char *image_path = "../model/cls/head/head.jpg";
    string label_path = "../model/cls/head/label.json";\
    // 默认的NMS阈值
    float nms_thresh = 0.25;
    // 默认的置信度阈值
    float box_thresh = 0.45;

    int class_num ;
    nlohmann::json label_json;
    read_json(label_path,class_num,label_json);
    int ret;
    // 初始化模型
    rknn_app_context_t rknn_app_ctx;
    memset(&rknn_app_ctx, 0, sizeof(rknn_app_context_t));

    ret = init_yolov8_model(model_path, &rknn_app_ctx);
    if (ret != 0)
    {
        printf("init_yolov8_model fail! ret=%d model_path=%s\n", ret, model_path);
        ret = release_yolov8_model(&rknn_app_ctx);
        if (ret != 0)
        {
            printf("release_yolov8_model fail! ret=%d\n", ret);
        }

    }

    // 读取图片
    image_buffer_t src_image;
    memset(&src_image, 0, sizeof(image_buffer_t));
    ret = read_image(image_path, &src_image);
    if (ret != 0)
    {
        printf("read image fail! ret=%d image_path=%s\n", ret, image_path);
        ret = release_yolov8_model(&rknn_app_ctx);
        if (ret != 0)
        {
            printf("release_yolov8_model fail! ret=%d\n", ret);
        }

        if (src_image.virt_addr != NULL)
        {
            free(src_image.virt_addr);
        }
    }

    // 图片前处理
    rknn_input inputs[rknn_app_ctx.io_num.n_input];
    memset(inputs, 0, sizeof(inputs));
    letterbox_t letter_box;
    memset(&letter_box, 0, sizeof(letterbox_t));
    ret = pre_process(&rknn_app_ctx,&src_image,inputs,&letter_box);
    if (ret != 0)
    {
        printf("pre_process fail! ret=%d\n", ret);
        return -1;
    }

    // 推理
    rknn_output outputs[rknn_app_ctx.io_num.n_output];
    memset(outputs, 0, sizeof(outputs));
    ret = inference_yolov8_model(&rknn_app_ctx, inputs, outputs);
    if (ret != 0)
    {
        printf("init_yolov8_model fail! ret=%d\n", ret);
        ret = release_yolov8_model(&rknn_app_ctx);
        if (ret != 0)
        {
            printf("release_yolov8_model fail! ret=%d\n", ret);
        }

        if (src_image.virt_addr != NULL)
        {
            free(src_image.virt_addr);
        }
    }

    // 后处理
    object_detect_result od_results;
    if (!&od_results)
    {
        return -1;
    }
    memset(&od_results, 0x00, sizeof(&od_results));
    post_process(&rknn_app_ctx, outputs, box_thresh, &od_results, class_num);

    // Remeber to release rknn output
    rknn_outputs_release(rknn_app_ctx.rknn_ctx, rknn_app_ctx.io_num.n_output, outputs);

    // 画框和概率
    char text[256];

    object_detect_result *cls_result = &(od_results);
    const char* label = label_json[std::to_string(cls_result->id)].get<std::string>().c_str();
    sprintf(text, "%s %.1f%%", label, cls_result->score * 100);
    std::cout << "分类结果: " << text << std::endl; 
    draw_text(&src_image, text, 10, 20 - 20, COLOR_RED, 6);

    write_image("out.png", &src_image);
    return 0;
}
